Single grouped weight fixes by CarlosGomes98 · Pull Request #3225 · NVIDIA/TransformerEngine

CarlosGomes98 · 2026-07-21T09:57:53Z

Description

A few fixes to the single grouped weight implementation

12406c9 — Preserve grouped weight initialization metadata
99b3f37 — Mark late grouped weights for delayed wgrad
aa4e233 — Preserve grouped MXFP8 columnwise usage

See #3178 for the original PR.
PR #3224 is very relevant to this one, both are required fixes.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Change A
Change B

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

greptile-apps · 2026-07-21T10:04:20Z

Greptile Summary

This PR delivers three targeted bug fixes to the single-grouped-weight path introduced in #3178: it preserves the per-weight high-precision initialization values when packing discrete weights into a GroupedTensor, correctly marks late-attached grouped weights with skip_backward_post_hook for delayed wgrad, and stops set_usage from clearing already-set MXFP8 columnwise usage during eval-only forward passes.

base.py / grouped_linear.py: Factored out _get/clear/attach_high_precision_init_val helpers and used them in make_grouped_weights to stack and transfer per-weight FP32 master init values onto the new grouped parameter (raising RuntimeError on partial state).
ops/basic/grouped_linear.py: Overrides register_parameter to handle the meta-device shell case — marks the grouped weight with skip_backward_post_hook=True when it is attached post-construction, and guards _apply_delay_wgrad_param_hooks against a missing weight.
ops/fused/grouped_mlp.py: Changes set_usage(columnwise=input_requires_grad) to set_usage(columnwise=True if ... else None) so eval passes (where grad is not required) pass None rather than False, leaving the previously set columnwise flag intact.

Confidence Score: 5/5

The three fixes are well-scoped and each has a dedicated regression test; no existing behaviour is altered for the non-single-grouped-weight paths.

All three fixes are narrowly targeted at the single-grouped-weight path, the logic is straightforward, and each change is covered by a new test. The set_usage(columnwise=None) change correctly exploits the existing guard in the quantizer. The _high_precision_init_val stacking is safe and cleans up after itself. The register_parameter override guards with getattr for pre-init access. No regressions to unrelated paths.

No files require special attention; the changes are localized and well-tested.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Refactors inline get/clear closures into three importable module-level helpers. Behavior is functionally identical to the original; the helpers are now usable by grouped_linear.py.
transformer_engine/pytorch/module/grouped_linear.py	Adds high-precision init-val preservation in make_grouped_weights: reads per-weight values, validates all-or-none consistency, stacks them onto the new grouped parameter, and clears the originals.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Overrides register_parameter to mark a late-attached grouped weight with skip_backward_post_hook=True when delay_wgrad_compute is active. Also guards _apply_delay_wgrad_param_hooks against a missing weight on meta-device shells.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Fixes set_usage calls for single-grouped-weight FC1 and FC2 to pass columnwise=None (not False) on non-grad eval passes, preventing already-set columnwise usage from being cleared.
tests/pytorch/test_grouped_mlp.py	Adds two new targeted tests: one for the meta-device delayed-wgrad marking path and one for MXFP8 columnwise-usage preservation across train/eval forward passes.
tests/pytorch/test_sanity.py	Adds two sanity tests covering the high-precision init-val transfer: one checking numerical equivalence between discrete and grouped weight initialization, and one verifying the partial-init-val error is raised.

_{Reviews (2): Last reviewed commit: "Fix bug where grouped MLP used same quan..." | Re-trigger Greptile}

zhongbozhu · 2026-07-21T23:38:22Z

/te-ci pytorch L1

zhongbozhu

LGTM

phu0ngng · 2026-07-22T15:30:03Z

/te-ci pytorch L1

timmoon10

LGTM

timmoon10 · 2026-07-23T17:58:41Z

+                # Full-iteration CUDA graph replay does not rerun this Python metadata
+                # update. Remember that a grad-enabled forward requested columnwise
+                # storage so eager eval cannot drop buffers still used by captured dgrad.
+                fc1_weight_quantizer.set_usage(
+                    rowwise=True,
+                    columnwise=input_requires_grad or fc1_weight_quantizer.columnwise_usage,
+                )
                fc1_op.weight.quantizer = fc1_weight_quantizer
                grouped_fc1_weight = fc1_op.weight


The op and the param purposely have different quantizers so that you can change them independently. However, we are incorrectly reusing the same quantizer in two places, and getting unexpected interactions. The op's quantizer is set in each forward pass:

TransformerEngine/transformer_engine/pytorch/ops/basic/grouped_linear.py

Line 697 in aa4e233

weight_quantizer.set_usage(rowwise=True, columnwise=requires_grad)

The proper fix is keep the op and param quantizers distinct.

…d weight param Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 · 2026-07-23T18:22:02Z

/te-ci pytorch

CarlosGomes98 added 3 commits July 21, 2026 11:11

[PyTorch] Preserve grouped weight initialization metadata

12406c9

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

[PyTorch] Mark late grouped weights for delayed wgrad

99b3f37

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

[PyTorch] Preserve grouped MXFP8 columnwise usage

aa4e233

Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Jul 21, 2026

CarlosGomes98 marked this pull request as ready for review July 21, 2026 09:58

CarlosGomes98 requested review from ksivaman and timmoon10 as code owners July 21, 2026 09:58

CarlosGomes98 mentioned this pull request Jul 21, 2026

grouped linear single param fixes #3178

Closed

13 tasks

zhongbozhu approved these changes Jul 21, 2026

View reviewed changes

timmoon10 previously approved these changes Jul 23, 2026

View reviewed changes

timmoon10 reviewed Jul 23, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/fused/grouped_mlp.py Outdated

Comment thread transformer_engine/pytorch/ops/fused/grouped_mlp.py Outdated

Fix bug where grouped MLP used same quantizer in grouped linear op an…

0e91ea5

…d weight param Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 dismissed their stale review via 0e91ea5 July 23, 2026 18:21

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Single grouped weight fixes#3225

Single grouped weight fixes#3225
CarlosGomes98 wants to merge 4 commits into
NVIDIA:mainfrom
CarlosGomes98:cgomes/single-grouped-weight-independent-fixes

CarlosGomes98 commented Jul 21, 2026

Uh oh!

greptile-apps Bot commented Jul 21, 2026 •

edited

Loading

Uh oh!

zhongbozhu commented Jul 21, 2026

Uh oh!

zhongbozhu left a comment

Uh oh!

phu0ngng commented Jul 22, 2026

Uh oh!

timmoon10 left a comment

Uh oh!

timmoon10 Jul 23, 2026

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Jul 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

CarlosGomes98 commented Jul 21, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Uh oh!

zhongbozhu commented Jul 21, 2026

Uh oh!

zhongbozhu left a comment

Choose a reason for hiding this comment

Uh oh!

phu0ngng commented Jul 22, 2026

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

timmoon10 Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Jul 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

greptile-apps Bot commented Jul 21, 2026 •

edited

Loading